# 使用链式连接数据库


import os
from operator import itemgetter

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_deepseek import ChatDeepSeek

# 本地有clash代理，所以配置一下，不然一些依赖下载不下来（chroma）,或者将代理关闭后下载
os.environ['http_proxy'] = '127.0.0.1:7890'
os.environ['https_proxy'] = '127.0.0.1:7890'

os.environ["LANGCHAIN_TRACING_V3"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["LANGCHAIN_PROJECT"] = "langchain-community-demo"

deepseek_api_key = "sk-1dd16a258a73428d910d38c782e1c94f"

# 1、创建模型
model = ChatDeepSeek(
    api_key=deepseek_api_key,
    # deepseek-reasoner : DeepSeek-R1
    # deepseek-chat : DeepSeek-V3
    model="deepseek-reasoner",
    temperature=0.7,
    max_tokens=1024
)

# 本例以 SQLAlchemy 演示
HOSTNAME = "127.0.0.1"
PORT = "3306"
DATABASE = "book_system"
USERNAME = "root"
PASSWORD = "123456"
# mysqlclient驱动URL
# MYSQL_URL = f"mysql+mysqldb://{USERNAME}:{PASSWORD}@{HOSTNAME}:{PORT}/{DATABASE}?charset=utf8mb4"
MYSQL_URL = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URL)

# 测试连接
# all_effect_tb_names = db.get_usable_table_names()
# print(all_effect_tb_names)
# query_result = db.run("select * from sys_user_role")
# print(query_result)

# 方式一 大模型直接和数据库整合(只能生成SQL语句，不推荐)
llm_sql_chain = create_sql_query_chain(model, db)
# resp = llm_sql_chain.invoke({"question": "查询用户角色表有多少条数据？"})
# # 对于不同的模型，根据实际情况判断，deepseek、gpt-4-turbo返回SQL语句，gpt-3.5-turbo会直接返回SQL
# print(resp)

# 方式二
answer_prompt = PromptTemplate.from_template(
    """
    给定以下用户问题、SQL语句和SQL执行后的结果，回答用户问题。
    Question: {question}
    SQL Query: {query}
    SQL Result: {result}
    回答:
    """
)


"""注意：
不同的模型输出SQL结果格式不一样，下面去除前缀方式也不一样，此例中是deepseek-chat
"""
# 自定义一个函数来去除大模型生成的SQL前缀
def remove_sql_query_prefix(query):
    if query.startswith("SQLQuery:"):
        query = query[len("SQLQuery:"):].strip()
    query = query.replace("```sql", "").replace("```", "").strip()
    return query


# 创建一个执行SQL语句的工具
execute_sql_tool = QuerySQLDataBaseTool(db=db)
# 1、生成SQL 2、执行SQL
# assign中的key要和answer_prompt中的对应
chain = (RunnablePassthrough.assign(query=llm_sql_chain)
         .assign(clean_query=lambda inputs: remove_sql_query_prefix(inputs['query']))
         .assign(result=itemgetter('clean_query') | execute_sql_tool)
         | answer_prompt
         | model
         | StrOutputParser()
         )
resp = chain.invoke({"question": "查询用户角色表有多少条数据？"})
print(resp)
